package com.ruoyi.chat.utils;

import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

/**
 * Langchain4j工具类
 *
 * @author ruoyi
 */
@Component
public class LangchainUtils {
    private static final Logger log = LoggerFactory.getLogger(LangchainUtils.class);

    // Redis连接信息,默认使用本机
    @Value("127.0.0.1")
    private String redisHost;
    // Redis端口，默认6379
    @Value("6379")
    private Integer redisPort;

    // Embedding模型配置
    @Value("${langchain4j.open-ai.embedding-model.api-key}")
    private String embeddingApiKey;

    @Value("${langchain4j.open-ai.embedding-model.base-url}")
    private String embeddingBaseUrl;

    @Value("${langchain4j.open-ai.embedding-model.model-name}")
    private String embeddingModelName;

    @Value("${langchain4j.open-ai.embedding-model.timeout:60}")
    private Integer embeddingTimeout;

    @Value("${langchain4j.open-ai.embedding-model.dimensions:1536}")
    private Integer embeddingDimensions;

    /**
     * 获取Embedding模型
     *
     * @return EmbeddingModel Embedding模型
     */
    public EmbeddingModel getEmbeddingModel() {
        return OpenAiEmbeddingModel.builder()
                .apiKey(embeddingApiKey)
                .baseUrl(embeddingBaseUrl)
                .modelName(embeddingModelName)
                .timeout(java.time.Duration.ofSeconds(embeddingTimeout))
                .build();
    }

    /**
     * 获取Redis向量存储
     *
     * @param namespace 命名空间，用于区分不同的知识库
     * @return EmbeddingStore Redis向量存储
     */
    public EmbeddingStore<TextSegment> getLocalEmbeddedStore(String namespace) {
        try {
            log.info("创建Redis向量存储，主机：{}，端口：{}，维度：{}", redisHost, redisPort, embeddingDimensions);
            return RedisEmbeddingStore
                    .builder()
                    .host(redisHost)
                    .port(redisPort)
                    .dimension(embeddingDimensions)
                    .build();
        } catch (Exception e) {
            log.error("创建Redis向量存储失败", e);
            // 如果Redis配置有问题，直接报错！
            return null;
        }
    }

    /**
     * 查询与问题相关的文本片段
     *
     * @param question     用户问题
     * @param namespace    命名空间
     * @param maxResults   最大返回结果数
     * @param minRelevance 最小相关度
     * @return 相关文本片段列表
     */
    public List<String> findRelevantTexts(String question, String namespace, int maxResults, double minRelevance) {
        try {
            log.info("开始检索问题相关信息: {}, 命名空间: {}, 最大结果数: {}, 最小相关度: {}",
                    question, namespace, maxResults, minRelevance);

            // 获取embedding模型和存储
            EmbeddingModel embeddingModel = getEmbeddingModel();
            EmbeddingStore<TextSegment> embeddingStore = getLocalEmbeddedStore(namespace);

            // 对问题进行向量化
            Embedding questionEmbedding = embeddingModel.embed(question).content();

            // 创建搜索请求
            EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                    .queryEmbedding(questionEmbedding)
                    .maxResults(maxResults)
                    .minScore(minRelevance)
                    .build();

            // 执行相似度搜索
            EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(searchRequest);
            List<EmbeddingMatch<TextSegment>> relevantSegments = searchResult.matches();

            log.info("检索到相关片段数量: {}", relevantSegments.size());

            // 转换为文本列表
            List<String> relevantTexts = relevantSegments.stream()
                    .map(match -> {
                        TextSegment segment = match.embedded();
                        double relevance = match.score();
                        log.debug("相关片段 (相关度: {}): {}", relevance, segment.text());
                        return segment.text();
                    })
                    .collect(Collectors.toList());

            return relevantTexts;
        } catch (Exception e) {
            log.error("执行相关文本检索时出错", e);
            return Collections.emptyList();
        }
    }

    /**
     * 在所有知识库中检索相关信息
     *
     * @param question     用户问题
     * @param maxResults   每个知识库最大返回结果数
     * @param minRelevance 最小相关度
     * @return 相关文本片段列表
     */
    public List<String> searchAllKnowledgeBases(String question, int maxResults, double minRelevance) {
        log.info("在所有知识库中搜索: {}", question);

        List<String> allResults = new ArrayList<>();

        try {
            // 获取embedding模型
            EmbeddingModel embeddingModel = getEmbeddingModel();
            if (embeddingModel == null) {
                log.error("Embedding模型初始化失败");
                return allResults;
            }

            // 对问题进行向量化
            log.info("步骤2: 问题向量化处理");
            long embedStartTime = System.currentTimeMillis();
            Embedding questionEmbedding = embeddingModel.embed(question).content();
            long embedEndTime = System.currentTimeMillis();

            log.info("问题向量化完成，维度: {}, 耗时: {}ms",
                    questionEmbedding.vector().length,
                    (embedEndTime - embedStartTime));

            // 直接全局检索
            EmbeddingStore<TextSegment> embeddingStore = getLocalEmbeddedStore("embedding");
            EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                    .queryEmbedding(questionEmbedding)
                    .maxResults(maxResults)
                    .minScore(minRelevance)
                    .build();
            EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(searchRequest);
            List<EmbeddingMatch<TextSegment>> relevantSegments = searchResult.matches();
            // 后续处理...
            // 日志输出检索结果
            log.info("全局检索结果:");
            for (int i = 0; i < relevantSegments.size(); i++) {
                EmbeddingMatch<TextSegment> match = relevantSegments.get(i);
                TextSegment segment = match.embedded();
                double relevance = match.score();
                allResults.add(segment.text());
                log.info("片段 {} 相关度: {}", segment.text(), relevance);
            }
            log.info("========== 多知识库RAG检索流程结束 ==========");
        } catch (Exception e) {
            log.info("出错了");
        }
        return allResults;
    }
}
